import json
import random

import numpy as np
import os

from sklearn.model_selection import KFold, train_test_split
import pandas as pd

if __name__ == '__main__':
    with open('config.json', 'r') as f:
        config = json.load(f)

    mod = config["mod"]
    dataset_path = config["data"]["dataset_path"]
    data_path = dataset_path + config["data"]["augmentation_data_path"] + mod + '/'
    output_path = config["list"]["list_path"] + mod + '/'
    num_augmentations = config["num_augmentations"]
    train_size = config["train_size"]

    num_samples = config["data"]["num_data"]
    sample_list = list(range(1, num_samples + 1))
    sample_name = 'A{0}_{1}{2}.ply'

    if not os.path.exists(output_path):
        os.mkdir(output_path)

    # get valid sample list
    valid_sample_list = []
    for i_sample in sample_list:
        for i_aug in range(num_augmentations):
            if os.path.exists(os.path.join(data_path, sample_name.format(i_aug, mod, i_sample))):
                valid_sample_list.append(i_sample)

    # remove duplicated
    sample_list = list(dict.fromkeys(valid_sample_list))
    random.shuffle(sample_list)
    sample_list = (np.asarray(sample_list))

    i_cv = 0
    kf = KFold(n_splits=config["list"]["n_split"], shuffle=False)
    for train_idx, test_idx in kf.split(sample_list):

        i_cv += 1
        print('Round:', i_cv)

        train_list, test_list = sample_list[train_idx], sample_list[test_idx]
        train_list, val_list = train_test_split(train_list, train_size=0.8, shuffle=True)

        print('Training list:\n', train_list,
              '\nValidation list:\n', val_list,
              '\nTest list:\n', test_list)

        # training
        train_name_list = []
        for i_sample in train_list:
            for i_aug in range(num_augmentations):
                subject_name = 'A{}_{}{}.ply'.format(i_aug, mod, i_sample)
                train_name_list.append(os.path.join(data_path, subject_name))

        with open(os.path.join(output_path, 'train_list_{}.csv'.format(i_cv)), 'w') as file:
            for f in train_name_list:
                file.write(f+'\n')

        # validation
        val_name_list = []
        for i_sample in val_list:
            for i_aug in range(num_augmentations):
                subject_name = 'A{}_{}{}.ply'.format(i_aug, mod, i_sample)
                val_name_list.append(os.path.join(data_path, subject_name))

        with open(os.path.join(output_path, 'val_list_{}.csv'.format(i_cv)), 'w') as file:
            for f in val_name_list:
                file.write(f+'\n')

        # test
        test_df = pd.DataFrame(data=test_list, columns=['Test ID'])
        test_df.to_csv(os.path.join(output_path, 'test_list_{}.csv'.format(i_cv)), index=False)

        print('--------------------------------------------')
        print('# of train:', len(train_name_list))
        print('# of validation:', len(val_name_list))
        print('--------------------------------------------')